/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once

#include <baja/common.hpp>
#include <baja/heap.hpp>
#include <map>
#include <vector>
#include <limits>
#include <functional>

namespace baja {

template <class K, class T>
class mru_cache
{
private:
    struct heap_item
    {
        K key;
        size_t touch;
    };

    typedef typename heap<heap_item>::handle heap_handle;

    struct map_item
    {
        K key;
        T data;
        heap_handle h;
    };

public:
    mru_cache(const size_t capacity) :
        _heap(capacity, std::bind(&mru_cache<K, T>::compare, std::placeholders::_1, std::placeholders::_2)),
        //_purgeHandler(std::bind(&mru_cache<K, T>::purge, std::placeholders::_1, std::placeholders::_2)),
        _lastTouch(0)
    {
    }

    void setPurgeHandler(std::function<void(const K&, T&)> handler)
    {
        _purgeHandler = handler;
    }

    T operator[](const K& key)
    {
        return get(key);
    }

    bool contains(const K& key)
    {
        auto it = _map.find(key);
        return (it != _map.end());
    }

    void add(const K& key, const T& data, bool failIfExists = false)
    {
        if (!tryAdd(key, data, failIfExists)) throw std::invalid_argument("item already in the cache");
    }

    bool tryAdd(const K& key, const T& data, bool failIfExists = false)
    {
        auto it = _map.find(key);
        if (it != _map.end())
        {
            // already in the cache
            if (!failIfExists)
            {
                // update the value and touch
                it->second.data = data;
                this->touch(it->second.h);
            }

            return false; // already in the cache
        }

        // doesn't exist, so add it
        if (_heap.full())
        {
            // if the heap is full make room
            heap_item hi = _heap.pop();
            _purgeHandler(hi.key, _map[hi.key].data);
            _map.erase(hi.key);
        }

        // add to heap and map
        heap_item hi;
        hi.key = key;
        hi.touch = ++_lastTouch;

        map_item mi;
        mi.key = key;
        mi.data = data;
        mi.h = _heap.push(hi);
        _map[key] = mi;

        return true;
    }

    T get(const K& key)
    {
        T t;
        if (!tryGet(key, &t)) throw std::invalid_argument("item doesn't exist in the cache");
        return t;
    }

    bool tryGet(const K& key, T* data)
    {
        auto it = _map.find(key);
        if (it == _map.end())
        {
            return false;
        }
        
        map_item& mi = it->second;
        *data = mi.data;
        touch(mi.h);
        return true;
    }

    void remove(const K& key)
    {
        auto it = _map.find(key);
        if (it == _map.end())
        {
            return;
        }

        _heap.remove(it->second.h);
        _purgeHandler(key, it->second.data);
        _map.erase(it);
    }

    bool full() { return _heap.full(); }
    bool empty() { return _heap.empty(); }
    size_t size() { return _heap.size(); }

    void clear()
    {
        std::vector<K> keys;

        enumerate(
            [&](const K& key, T& data)
            {
                keys.push_back(key);
            });

        std::for_each(
            keys.begin(), keys.end(),
            [this](const K& key)
            {
                remove(key);
            });
    }

    void enumerate(std::function<void(const K&, T&)> f)
    {
        std::for_each(
            _map.begin(), _map.end(),
            [f](std::pair<K, map_item> pair)
            {
                f(pair.second.key, pair.second.data);
            });
    }

private:
    void touch(const heap_handle h)
    {
        _heap.touch(h, [this](heap_item& hi) { hi.touch = ++_lastTouch; });
    }

    static int compare(const heap_item& left, const heap_item& right)
    {
        return left.touch > right.touch;
    }
    
    static void purge(const K& key, T& data)
    {
        // noop
    }

    std::map<K, map_item> _map;
    heap<heap_item> _heap;
    size_t _lastTouch;
    std::function<void(const K&, T&)> _purgeHandler;
};

} // namespace baja